﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using Microsoft.Cci;
using Microsoft.Cci.MetadataReader;
using Microsoft.Cci.MutableCodeModel;
using System.Runtime.Serialization; // needed for defining exception .ctors

namespace HeapProfiler
{
	/// <summary>
	/// A mutator that modifies method bodies at the IL level.
	/// It injects a call to ShadowMemory for each store
	/// to a local/static for which the PDB reader is able to provide a name.
	/// </summary>
	public class ILRewrite : MetadataMutator
	{
		private PdbReader pdbReader;

		private IMethodReference fref_initstate;
		private IMethodReference fref_initithparam;

		private IMethodReference fref_profileatcallentry;

		private String m_specificfname;

		private int m_methodsampletoken = 1;

		/// <summary>
		/// Returns true if the type is a memory based reference (we count structs as memory types).
		/// </summary>
		public static bool IsMemoryType(ITypeReference type)
		{
			bool isPrimitive = type.TypeCode != PrimitiveTypeCode.NotPrimitive;
			bool isString = type.TypeCode == PrimitiveTypeCode.String;

			return (!isPrimitive) || isString; //for some reason string is primitive?
		}

		public ILRewrite(IMetadataHost host, PdbReader pdbReader, IAssembly runtimeSupport, String extrafname)
			: base(host, true)
		{
			this.pdbReader = pdbReader;

			var nameTable = host.NameTable;
			var platformType = host.PlatformType;
			var systemInt = platformType.SystemInt32;
			var systemString = platformType.SystemString;
			var systemVoid = platformType.SystemVoid;
			var objectType = platformType.SystemObject;
			var systemType = platformType.SystemType;

			foreach(ITypeDefinition td in runtimeSupport.GetAllTypes())
			{
				String ftn = td.ToString();
				if(ftn.Equals("HeapProfiler.Runtime.HeapWalk"))
				{
					foreach(IMethodDefinition md in td.Methods)
					{
						if(md.Name.Value.Equals("InitDumpState"))
							this.fref_initstate = md;
						else if(md.Name.Value.Equals("SetIthParamInfo"))
							this.fref_initithparam = md;
						else if(md.Name.Value.Equals("SampleMemProfileAtCallEntry"))
							this.fref_profileatcallentry = md;
						else
						{ ;}
					}
				}
			}

			this.m_specificfname = extrafname;

			this.m_methodsampletoken = 1;
		}

		private static String[] badnamespaces = { "System.", "Microsoft.", "Windows.", "<PrivateImplementationDetails>" };

		public override MethodBody Visit(MethodBody methodBody)
		{
			IMethodDefinition currentMethod = this.GetCurrentMethod();
			methodBody.MethodDefinition = currentMethod;

			String classname = currentMethod.ContainingType.ToString();
			bool usingspecificname = (this.m_specificfname != null);
			bool specificnameok = usingspecificname && (this.m_specificfname.Equals(currentMethod.Name.Value));

			bool isgettersetter = (currentMethod.Name.Value.StartsWith("get_") | currentMethod.Name.Value.StartsWith("set_"));
			bool isprivate = (currentMethod.Visibility == TypeMemberVisibility.Private);
			bool isstatic = (methodBody.MethodDefinition.IsStatic);
			bool isconstructor = (methodBody.MethodDefinition.IsConstructor || methodBody.MethodDefinition.IsStaticConstructor);
			bool isabstract = methodBody.MethodDefinition.IsAbstract;
			bool badnamespace = Array.Exists(badnamespaces, (String bns) => classname.StartsWith(bns));
			bool byrefarg = methodBody.MethodDefinition.Parameters.Any((IParameterDefinition pd) => pd.IsByReference | pd.IsOut | pd.IsParameterArray);
			bool byrefrcvr = methodBody.MethodDefinition.ContainingType.IsValueType;
			try
			{
				if(usingspecificname)
				{
					if(specificnameok)
						ProcessOperations(methodBody);
					else
						base.Visit(methodBody);
				}
				else
				{
					if((!badnamespace) & (!byrefrcvr) & (!byrefarg) & (!isstatic) & (!isconstructor) & (!isabstract) & (!isprivate) & (!isgettersetter))
						ProcessOperations(methodBody);
					else
						base.Visit(methodBody);
				}
			}
			catch(Exception e)
			{
				Console.WriteLine("Internal error during IL mutation for the method '{0}'." + e,
				  MemberHelper.GetMemberSignature(currentMethod, NameFormattingOptions.SmartTypeName));
			}

			return methodBody;
		}

		/// <summary>
		/// Rewrite the method to do the sampling for the memory profiler at call entry.
		/// </summary>
		private void ProcessOperations(MethodBody methodBody)
		{
			List<IOperation> operations = methodBody.Operations;
			int count = methodBody.Operations.Count;

			ILGenerator generator = new ILGenerator(this.host, this.GetCurrentMethod());

			var methodName = MemberHelper.GetMemberSignature(methodBody.MethodDefinition, NameFormattingOptions.SmartTypeName);

			////
			//Record all offsets that appear as part of an exception handler

			Dictionary<uint, bool> offsetsUsedInExceptionInformation = new Dictionary<uint, bool>();
			foreach(var exceptionInfo in methodBody.OperationExceptionInformation)
			{
				uint x = exceptionInfo.TryStartOffset;
				if(!offsetsUsedInExceptionInformation.ContainsKey(x)) offsetsUsedInExceptionInformation.Add(x, true);
				x = exceptionInfo.TryEndOffset;
				if(!offsetsUsedInExceptionInformation.ContainsKey(x)) offsetsUsedInExceptionInformation.Add(x, true);
				x = exceptionInfo.HandlerStartOffset;
				if(!offsetsUsedInExceptionInformation.ContainsKey(x)) offsetsUsedInExceptionInformation.Add(x, true);
				x = exceptionInfo.HandlerEndOffset;
				if(!offsetsUsedInExceptionInformation.ContainsKey(x)) offsetsUsedInExceptionInformation.Add(x, true);
				if(exceptionInfo.HandlerKind == HandlerKind.Filter)
				{
					x = exceptionInfo.FilterDecisionStartOffset;
					if(!offsetsUsedInExceptionInformation.ContainsKey(x)) offsetsUsedInExceptionInformation.Add(x, true);
				}
			}

			////
			//Pass 1: Make a label for each branch target

			Dictionary<uint, ILGeneratorLabel> offset2Label = new Dictionary<uint, ILGeneratorLabel>();
			for(int i = 0; i < count; i++)
			{
				IOperation op = operations[i];
				switch(op.OperationCode)
				{
					case OperationCode.Beq:
					case OperationCode.Bge:
					case OperationCode.Bge_Un:
					case OperationCode.Bgt:
					case OperationCode.Bgt_Un:
					case OperationCode.Ble:
					case OperationCode.Ble_Un:
					case OperationCode.Blt:
					case OperationCode.Blt_Un:
					case OperationCode.Bne_Un:
					case OperationCode.Br:
					case OperationCode.Brfalse:
					case OperationCode.Brtrue:
					case OperationCode.Leave:
					case OperationCode.Beq_S:
					case OperationCode.Bge_S:
					case OperationCode.Bge_Un_S:
					case OperationCode.Bgt_S:
					case OperationCode.Bgt_Un_S:
					case OperationCode.Ble_S:
					case OperationCode.Ble_Un_S:
					case OperationCode.Blt_S:
					case OperationCode.Blt_Un_S:
					case OperationCode.Bne_Un_S:
					case OperationCode.Br_S:
					case OperationCode.Brfalse_S:
					case OperationCode.Brtrue_S:
					case OperationCode.Leave_S:
						uint x = (uint)op.Value;
						if(!offset2Label.ContainsKey(x))
							offset2Label.Add(x, new ILGeneratorLabel());
						break;
					case OperationCode.Switch:
						uint[] offsets = op.Value as uint[];
						foreach(var offset in offsets)
						{
							if(!offset2Label.ContainsKey(offset))
								offset2Label.Add(offset, new ILGeneratorLabel());
						}
						break;
					default:
						break;
				}
			}

			////
			//Pass 2: Emit each operation, along with labels

			for(int i = 0; i < count; i++)
			{
				IOperation op = operations[i];
				ILGeneratorLabel label;
				if(op.Location is IILLocation)
				{
					generator.MarkSequencePoint(op.Location);
				}
				//Mark operation if it is a label for a branch
				if(offset2Label.TryGetValue(op.Offset, out label))
				{
					generator.MarkLabel(label);
				}

				//Mark operation if it is pointed to by an exception handler
				bool ignore;
				uint offset = op.Offset;
				if(offsetsUsedInExceptionInformation.TryGetValue(offset, out ignore))
				{
					foreach(var exceptionInfo in methodBody.OperationExceptionInformation)
					{
						if(offset == exceptionInfo.TryStartOffset)
							generator.BeginTryBody();

						// Never need to do anthing when offset == exceptionInfo.TryEndOffset because
						// we pick up an EndTryBody from the HandlerEndOffset below
						//  generator.EndTryBody();

						if(offset == exceptionInfo.HandlerStartOffset)
						{
							switch(exceptionInfo.HandlerKind)
							{
								case HandlerKind.Catch:
									generator.BeginCatchBlock(exceptionInfo.ExceptionType);
									break;
								case HandlerKind.Fault:
									generator.BeginFaultBlock();
									break;
								case HandlerKind.Filter:
									generator.BeginFilterBody();
									break;
								case HandlerKind.Finally:
									generator.BeginFinallyBlock();
									break;
							}
						}
						if(exceptionInfo.HandlerKind == HandlerKind.Filter && offset == exceptionInfo.FilterDecisionStartOffset)
						{
							generator.BeginFilterBlock();
						}
						if(offset == exceptionInfo.HandlerEndOffset)
							generator.EndTryBody();
					}
				}

				//First action in a method is to push the call frame.
				if(i == 0)
				{
					Console.WriteLine("Rewriting Method: " + methodBody.MethodDefinition.Name.ToString() + " For Memory Profiling.");
					EmitMemorySampleForCallEntry(generator, methodBody.MethodDefinition);
				}

				//Emit operation along with any injection
				switch(op.OperationCode)
				{
					case OperationCode.Beq:
					case OperationCode.Bge:
					case OperationCode.Bge_Un:
					case OperationCode.Bgt:
					case OperationCode.Bgt_Un:
					case OperationCode.Ble:
					case OperationCode.Ble_Un:
					case OperationCode.Blt:
					case OperationCode.Blt_Un:
					case OperationCode.Bne_Un:
					case OperationCode.Br:
					case OperationCode.Brfalse:
					case OperationCode.Brtrue:
					case OperationCode.Beq_S:
					case OperationCode.Bge_S:
					case OperationCode.Bge_Un_S:
					case OperationCode.Bgt_S:
					case OperationCode.Bgt_Un_S:
					case OperationCode.Ble_S:
					case OperationCode.Ble_Un_S:
					case OperationCode.Blt_S:
					case OperationCode.Blt_Un_S:
					case OperationCode.Bne_Un_S:
					case OperationCode.Br_S:
					case OperationCode.Brfalse_S:
					case OperationCode.Brtrue_S:
					case OperationCode.Leave:
					case OperationCode.Leave_S:
						generator.Emit(ILGenerator.LongVersionOf(op.OperationCode), offset2Label[(uint)op.Value]);
						break;
					case OperationCode.Switch:
						uint[] offsets = op.Value as uint[];
						ILGeneratorLabel[] labels = new ILGeneratorLabel[offsets.Length];
						for(int j = 0, n = offsets.Length; j < n; j++)
						{
							labels[j] = offset2Label[offsets[j]];
						}
						generator.Emit(OperationCode.Switch, labels);
						break;
					default:
						if(op.Value == null)
						{
							generator.Emit(op.OperationCode);
							break;
						}
						var typeCode = System.Convert.GetTypeCode(op.Value);
						switch(typeCode)
						{
							case TypeCode.Byte:
								generator.Emit(op.OperationCode, (byte)op.Value);
								break;
							case TypeCode.Double:
								generator.Emit(op.OperationCode, (double)op.Value);
								break;
							case TypeCode.Int16:
								generator.Emit(op.OperationCode, (short)op.Value);
								break;
							case TypeCode.Int32:
								generator.Emit(op.OperationCode, (int)op.Value);
								break;
							case TypeCode.Int64:
								generator.Emit(op.OperationCode, (long)op.Value);
								break;
							case TypeCode.Object:
								IFieldReference fieldReference = op.Value as IFieldReference;
								if(fieldReference != null)
								{
									generator.Emit(op.OperationCode, this.Visit(fieldReference));
									break;
								}
								ILocalDefinition localDefinition = op.Value as ILocalDefinition;
								if(localDefinition != null)
								{
									generator.Emit(op.OperationCode, localDefinition);
									break;
								}
								IMethodReference methodReference = op.Value as IMethodReference;
								if(methodReference != null)
								{
									generator.Emit(op.OperationCode, this.Visit(methodReference));
									break;
								}
								IParameterDefinition parameterDefinition = op.Value as IParameterDefinition;
								if(parameterDefinition != null)
								{
									generator.Emit(op.OperationCode, parameterDefinition);
									break;
								}
								ISignature signature = op.Value as ISignature;
								if(signature != null)
								{
									generator.Emit(op.OperationCode, signature);
									break;
								}
								ITypeReference typeReference = op.Value as ITypeReference;
								if(typeReference != null)
								{
									generator.Emit(op.OperationCode, this.Visit(typeReference));
									break;
								}
								throw new ILMutatorException("Should never get here: no other IOperation argument types should exist");
							case TypeCode.SByte:
								generator.Emit(op.OperationCode, (sbyte)op.Value);
								break;
							case TypeCode.Single:
								generator.Emit(op.OperationCode, (float)op.Value);
								break;
							case TypeCode.String:
								generator.Emit(op.OperationCode, (string)op.Value);
								break;
							default:
								// The other cases are the other enum values that TypeCode has.
								// But no other argument types should be in the Operations. ILGenerator cannot handle anything else,
								// so such IOperations should never exist.
								//case TypeCode.Boolean:
								//case TypeCode.Char:
								//case TypeCode.DateTime:
								//case TypeCode.DBNull:
								//case TypeCode.Decimal:
								//case TypeCode.Empty: // this would be the value for null, but the case when op.Value is null is handled before the switch statement
								//case TypeCode.UInt16:
								//case TypeCode.UInt32:
								//case TypeCode.UInt64:
								throw new ILMutatorException("Should never get here: no other IOperation argument types should exist");
						}
						break;
				}
			}

			////
			//Retrieve the operations (and the exception information) from the generator
			generator.AdjustBranchSizesToBestFit();
			methodBody.OperationExceptionInformation = new List<IOperationExceptionInformation>(generator.GetOperationExceptionInformation());
			methodBody.Operations = new List<IOperation>(generator.GetOperations());
			methodBody.MaxStack++;
			methodBody.MaxStack += 5;
		}

		/// <summary>
		/// Handle the emission of needed shadow operations for call entry.
		/// </summary>
		private void EmitMemorySampleForCallEntry(ILGenerator generator, IMethodDefinition mdef)
		{
			//initialize the dump state
			generator.Emit(OperationCode.Ldc_I4, mdef.ParameterCount + (mdef.IsStatic ? 0 : 1));
			generator.Emit(OperationCode.Call, this.fref_initstate);

			//store all the param info as needed

			if(!mdef.IsStatic)
			{
				generator.Emit(OperationCode.Ldc_I4_0);
				generator.Emit(OperationCode.Ldarg_0);
				generator.Emit(OperationCode.Ldc_I4_1);
				generator.Emit(OperationCode.Ldstr, "this");
				generator.Emit(OperationCode.Call, this.fref_initithparam);
			}

			int i = (mdef.IsStatic ? 0 : 1);
			foreach(IParameterDefinition pd in mdef.Parameters)
			{
				generator.Emit(OperationCode.Ldc_I4, i);

				if(IsMemoryType(pd.Type))
				{
					generator.Emit(OperationCode.Ldarg, pd);

					if(pd.Type.IsValueType)
						generator.Emit(OperationCode.Box, pd.Type);

					generator.Emit(OperationCode.Ldc_I4_1);
				}
				else
				{
					generator.Emit(OperationCode.Ldnull);
					generator.Emit(OperationCode.Ldc_I4_0);
				}

				generator.Emit(OperationCode.Ldstr, pd.Name.Value);
				generator.Emit(OperationCode.Call, this.fref_initithparam);

				++i;
			}

			//do the dump

			//push the samplename on the stack
			generator.Emit(OperationCode.Ldstr, mdef.ContainingTypeDefinition.ToString() + (mdef.IsStatic ? "#" : "@") + mdef.Name.ToString());
			generator.Emit(OperationCode.Ldc_I4, this.m_methodsampletoken++);
			generator.Emit(OperationCode.Call, this.fref_profileatcallentry);
		}

		/// <summary>
		/// Exceptions thrown during extraction. Should not escape this class.
		/// </summary>
		private class ILMutatorException : Exception
		{
			/// <summary>
			/// Exception specific to an error occurring in the contract extractor
			/// </summary>
			public ILMutatorException() { }
			/// <summary>
			/// Exception specific to an error occurring in the contract extractor
			/// </summary>
			public ILMutatorException(string s) : base(s) { }
			/// <summary>
			/// Exception specific to an error occurring in the contract extractor
			/// </summary>
			public ILMutatorException(string s, Exception inner) : base(s, inner) { }
			/// <summary>
			/// Exception specific to an error occurring in the contract extractor
			/// </summary>
			public ILMutatorException(SerializationInfo info, StreamingContext context) : base(info, context) { }
		}
	}
}
